{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Word Squares"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #trie #array #string #backtracking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #字典树 #数组 #字符串 #回溯"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: wordSquares"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #单词方块"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给定一个单词集合&nbsp;<code>words</code> <strong>（没有重复）</strong>，找出并返回其中所有的 <a href=\"https://en.wikipedia.org/wiki/Word_square\">单词方块</a><strong>&nbsp;</strong>。&nbsp;<code>words</code>&nbsp;中的同一个单词可以被 <strong>多次</strong> 使用。你可以按 <strong>任意顺序</strong> 返回答案。</p>\n",
    "\n",
    "<p>一个单词序列形成了一个有效的 <strong>单词方块</strong> 的意思是指从第 <code>k</code> 行和第 <code>k</code> 列 &nbsp;<code>0 &lt;= k &lt; max(numRows, numColumns)</code> 来看都是相同的字符串。</p>\n",
    "\n",
    "<ul>\n",
    "\t<li>例如，单词序列&nbsp;<code>[\"ball\",\"area\",\"lead\",\"lady\"]</code>&nbsp;形成了一个单词方块，因为每个单词从水平方向看和从竖直方向看都是相同的。</li>\n",
    "</ul>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>示例 1：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>words = [\"area\",\"lead\",\"wall\",\"lady\",\"ball\"]\n",
    "<strong>输出:</strong> [[\"ball\",\"area\",\"lead\",\"lady\"],\n",
    "[\"wall\",\"area\",\"lead\",\"lady\"]]\n",
    "<strong>解释：</strong>\n",
    "输出包含两个单词方块，输出的顺序不重要，只需要保证每个单词方块内的单词顺序正确即可。 \n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 2：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>words = [\"abat\",\"baba\",\"atan\",\"atal\"]\n",
    "<strong>输出：</strong>[[\"baba\",\"abat\",\"baba\",\"atal\"],\n",
    "[\"baba\",\"abat\",\"baba\",\"atan\"]]\n",
    "<strong>解释：</strong>\n",
    "输出包含两个单词方块，输出的顺序不重要，只需要保证每个单词方块内的单词顺序正确即可。 \n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示:</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>1 &lt;= words.length &lt;= 1000</code></li>\n",
    "\t<li><code>1 &lt;= words[i].length &lt;= 4</code></li>\n",
    "\t<li><code>words[i]</code>&nbsp;长度相同</li>\n",
    "\t<li><code>words[i]</code>&nbsp;只由小写英文字母组成</li>\n",
    "\t<li><code>words[i]</code>&nbsp;都 <strong>各不相同</strong></li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [word-squares](https://leetcode.cn/problems/word-squares/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [word-squares](https://leetcode.cn/problems/word-squares/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[\"area\",\"lead\",\"wall\",\"lady\",\"ball\"]', '[\"abat\",\"baba\",\"atan\",\"atal\"]']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "\n",
    "   def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "\n",
    "       self.words = words\n",
    "       self.N = len(words[0])\n",
    "       self.buildPrefixHashTable(self.words)\n",
    "\n",
    "       results = []\n",
    "       word_squares = []\n",
    "       for word in words:\n",
    "           word_squares = [word]\n",
    "           self.backtracking(1, word_squares, results)\n",
    "       return results\n",
    "\n",
    "   def backtracking(self, step, word_squares, results):\n",
    "       if step == self.N:\n",
    "           results.append(word_squares[:])\n",
    "           return\n",
    "\n",
    "       prefix = ''.join([word[step] for word in word_squares])\n",
    "       for candidate in self.getWordsWithPrefix(prefix):\n",
    "           word_squares.append(candidate)\n",
    "           self.backtracking(step+1, word_squares, results)\n",
    "           word_squares.pop()\n",
    "\n",
    "   def buildPrefixHashTable(self, words):\n",
    "       self.prefixHashTable = {}\n",
    "       for word in words:\n",
    "           for prefix in (word[:i] for i in range(1, len(word))):\n",
    "               self.prefixHashTable.setdefault(prefix, set()).add(word)\n",
    "\n",
    "   def getWordsWithPrefix(self, prefix):\n",
    "       if prefix in self.prefixHashTable:\n",
    "           return self.prefixHashTable[prefix]\n",
    "       else:\n",
    "           return set([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "\n",
    "        self.words = words\n",
    "        self.N = len(words[0])\n",
    "        self.buildPrefixHashTable(self.words)\n",
    "\n",
    "        results = []\n",
    "        word_squares = []\n",
    "        for word in words:\n",
    "            word_squares = [word]\n",
    "            self.backtracking(1, word_squares, results)\n",
    "        return results\n",
    "\n",
    "    def backtracking(self, step, word_squares, results):\n",
    "        if step == self.N:\n",
    "            results.append(word_squares[:])\n",
    "            return\n",
    "\n",
    "        prefix = ''.join([word[step] for word in word_squares])\n",
    "        for candidate in self.getWordsWithPrefix(prefix):\n",
    "            word_squares.append(candidate)\n",
    "            self.backtracking(step+1, word_squares, results)\n",
    "            word_squares.pop()\n",
    "\n",
    "    def buildPrefixHashTable(self, words):\n",
    "        self.prefixHashTable = {}\n",
    "        for word in words:\n",
    "            for prefix in (word[:i] for i in range(1, len(word))):\n",
    "                self.prefixHashTable.setdefault(prefix, set()).add(word)\n",
    "\n",
    "    def getWordsWithPrefix(self, prefix):\n",
    "        if prefix in self.prefixHashTable:\n",
    "            return self.prefixHashTable[prefix]\n",
    "        else:\n",
    "            return set([])\n",
    "'''\n",
    "作者：LeetCode\n",
    "链接：https://leetcode.cn/problems/word-squares/solutions/166154/dan-ci-fang-kuai-by-leetcode/\n",
    "来源：力扣（LeetCode）\n",
    "著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "        words_1 = {chr(i):[] for i in range(ord('a'), ord('a') + 26)}\n",
    "        for word in words:\n",
    "            words_1[word[0]].append(word)\n",
    "\n",
    "\n",
    "        ans = []\n",
    "        def isgood(now):\n",
    "            n = min(len(now), len(now[0]))\n",
    "            for i in range(n):\n",
    "                for j in range(i + 1, n):\n",
    "                    if now[i][j] != now[j][i]:\n",
    "                        return False\n",
    "            return True\n",
    "\n",
    "        def dfs(now):\n",
    "            if not isgood(now):\n",
    "                return\n",
    "            n, m = len(now), len(now[0])\n",
    "            if n == m:\n",
    "                ans.append(deepcopy(now))\n",
    "                return\n",
    "            for next in words_1[now[0][n]]:\n",
    "                now.append(next)\n",
    "                dfs(now)\n",
    "                now.pop()\n",
    "\n",
    "\n",
    "        for word in words:\n",
    "            dfs([word])\n",
    "        return ans\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "            \n",
    "\n",
    "            \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "        prefix_to_words = self.get_prefix_to_words(words)\n",
    "        result = []\n",
    "        for word in words:\n",
    "            self.dfs([word], result, prefix_to_words)\n",
    "        return result\n",
    "\n",
    "    def dfs(self, square, result, prefix_to_words):\n",
    "        nrow = len(square[0])\n",
    "        row_index = len(square)\n",
    "        if row_index == nrow:\n",
    "            result.append(square.copy())\n",
    "            return\n",
    "        prefix = \"\".join([square[i][row_index] for i in range(row_index)])\n",
    "        for word in prefix_to_words.get(prefix, []):\n",
    "            square.append(word)\n",
    "            self.dfs(square, result, prefix_to_words)\n",
    "            square.pop()\n",
    "\n",
    "    def get_prefix_to_words(self, words):\n",
    "        prefix_to_words = {}\n",
    "        for word in words:\n",
    "            for i in range(len(word)):\n",
    "                prefix = word[:i]\n",
    "                if prefix not in prefix_to_words:\n",
    "                    prefix_to_words[prefix] = set([word])\n",
    "                else:\n",
    "                    prefix_to_words[prefix].add(word)\n",
    "        return prefix_to_words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "\n",
    "\n",
    "class Trie:\n",
    "    def __init__(self):\n",
    "        self.children = [None for _ in range(26)]\n",
    "        self.index = -1\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "        # 所有的单词方块\n",
    "        root = Trie()\n",
    "\n",
    "        def add(i, w):\n",
    "            cur = root\n",
    "            for c in w:\n",
    "                index = ord(c) - ord('a')\n",
    "                if cur.children[index] == None:\n",
    "                    cur.children[index] = Trie()\n",
    "                cur = cur.children[index]\n",
    "            cur.index = i\n",
    "\n",
    "        for i, w in enumerate(words):\n",
    "            add(i, w)\n",
    "\n",
    "        def dfs(node, res):\n",
    "            if node.index != -1:\n",
    "                res.append(words[node.index])\n",
    "                return\n",
    "            for i in range(26):\n",
    "                if node.children[i]:\n",
    "                    dfs(node.children[i], res)\n",
    "\n",
    "        def search(data):\n",
    "            cur = root\n",
    "            for c in data:\n",
    "                index = ord(c) - ord('a')\n",
    "                if cur.children[index] is None:\n",
    "                    return []\n",
    "                cur = cur.children[index]\n",
    "            res = []\n",
    "            dfs(cur, res)\n",
    "            return res\n",
    "\n",
    "        n = len(words[0])\n",
    "        ans = []\n",
    "\n",
    "        def dfs1(index, ele):\n",
    "            if index == n:\n",
    "                ans.append(ele[:])\n",
    "                return\n",
    "            # 处理index行 需要之前的每一个index列位置的字符\n",
    "            cur = [w[index] for w in ele]\n",
    "            nex = search(cur)\n",
    "            for next_word in nex:\n",
    "                ele.append(next_word)\n",
    "                dfs1(index + 1, ele)\n",
    "                ele.pop()\n",
    "\n",
    "        for w in words:\n",
    "            dfs1(1, [w])\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Trie(object):\n",
    "    def __init__(self): \n",
    "        self.words = []\n",
    "        self.childs = collections.defaultdict(Trie)\n",
    "\n",
    "class Solution:\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "\n",
    "        root = Trie()\n",
    "        for word in words:\n",
    "            cur = root\n",
    "            for ch in word:                \n",
    "                cur = cur.childs[ch]\n",
    "                cur.words.append(word)\n",
    "        \n",
    "        def findWordsInPreifx(prefix: str) -> List[str]:\n",
    "            cur, found = root, True\n",
    "            for ch in prefix:\n",
    "                if ch not in cur.childs:\n",
    "                    found = False\n",
    "                    break\n",
    "                cur = cur.childs[ch]\n",
    "            return cur.words if found else []\n",
    "\n",
    "        res = []\n",
    "        def backtrack(square):\n",
    "            if len(square) == len(square[0]): \n",
    "                res.append(square[:])\n",
    "                return\n",
    "\n",
    "            columnInd = len(square)\n",
    "            prefix = [row[columnInd] for row in square]\n",
    "\n",
    "            for nxt in findWordsInPreifx(prefix):\n",
    "                square.append(nxt)\n",
    "                backtrack(square)\n",
    "                square.pop()\n",
    "\n",
    "        for word in words:\n",
    "            backtrack([word])\n",
    "        return res \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "        prefix_to_word = self.get_prefix_to_words(words)\n",
    "        result = []\n",
    "        path = []\n",
    "        visited = set()\n",
    "        for word in words:\n",
    "            self.dfs(word, result, path, visited, prefix_to_word)\n",
    "        return result\n",
    "\n",
    "    def dfs(self, word, result, path, visited, prefix_to_word):\n",
    "        path.append(word)\n",
    "        if len(path) == len(path[0]):\n",
    "            result.append(path.copy())\n",
    "            path.pop()\n",
    "            return\n",
    "\n",
    "        index = len(path)\n",
    "        for i in range(index, len(word)):\n",
    "            prefix = \"\".join([path[j][i] for j in range(index)])\n",
    "        if prefix not in prefix_to_word:\n",
    "            path.pop()        \n",
    "            return\n",
    "\n",
    "        prefix = \"\".join([path[i][index] for i in range(index)])\n",
    "        for next_word in prefix_to_word.get(prefix, []):\n",
    "            self.dfs(next_word, result, path, visited, prefix_to_word)\n",
    "        path.pop()\n",
    "\n",
    "    def get_prefix_to_words(self, words):\n",
    "        prefix_to_words = {}\n",
    "        for word in words:\n",
    "            for i in range(len(word)):\n",
    "                prefix = word[:i]\n",
    "                if prefix not in prefix_to_words:\n",
    "                    prefix_to_words[prefix] = set([word])\n",
    "                else:\n",
    "                    prefix_to_words[prefix].add(word)\n",
    "        return prefix_to_words\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "        prefix_to_words = self.get_prefix_to_words(words)\n",
    "        result = []\n",
    "        square = []\n",
    "        for word in words:\n",
    "            self.dfs(word, square, result, prefix_to_words)\n",
    "        return result\n",
    "\n",
    "    def dfs(self, word, square, result, prefix_to_words):\n",
    "        square.append(word)\n",
    "        nrow = len(square[0])\n",
    "        row_index = len(square)\n",
    "        if row_index == nrow:\n",
    "            result.append(square.copy())\n",
    "            square.pop()\n",
    "            return\n",
    "        \n",
    "        prefix = \"\".join([square[i][row_index] for i in range(row_index)])\n",
    "        # 只添加有共同前缀的单词\n",
    "        for word in prefix_to_words.get(prefix, []):\n",
    "            # square.append(word)\n",
    "            self.dfs(word, square, result, prefix_to_words)\n",
    "            # square.pop()\n",
    "        square.pop()\n",
    "\n",
    "    def get_prefix_to_words(self, words):\n",
    "        prefix_to_words = {}\n",
    "        for word in words:\n",
    "            for i in range(len(word)):\n",
    "                prefix = word[:i]\n",
    "                if prefix not in prefix_to_words:\n",
    "                    prefix_to_words[prefix] = set([word])\n",
    "                else:\n",
    "                    prefix_to_words[prefix].add(word)\n",
    "        return prefix_to_words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "        prefix_to_words = self.get_prefix_to_words(words)\n",
    "        result = []\n",
    "        square = []\n",
    "        for word in words:\n",
    "            self.dfs(word, square, result, prefix_to_words)\n",
    "        return result\n",
    "\n",
    "    def dfs(self, word, square, result, prefix_to_words):\n",
    "        square.append(word)\n",
    "        nrow = len(square[0])\n",
    "        row_index = len(square)\n",
    "        if row_index == nrow:\n",
    "            result.append(square.copy())\n",
    "            square.pop()\n",
    "            return\n",
    "        \n",
    "        for i in range(row_index, nrow):\n",
    "            prefix = \"\".join([square[r][i] for r in range(row_index)])\n",
    "            if prefix not in prefix_to_words:\n",
    "                square.pop()\n",
    "                return\n",
    "\n",
    "        prefix = \"\".join([square[i][row_index] for i in range(row_index)])\n",
    "        # 只添加有共同前缀的单词\n",
    "        for word in prefix_to_words.get(prefix, []):\n",
    "            # square.append(word)\n",
    "            self.dfs(word, square, result, prefix_to_words)\n",
    "            # square.pop()\n",
    "        square.pop()\n",
    "\n",
    "    def get_prefix_to_words(self, words):\n",
    "        prefix_to_words = {}\n",
    "        for word in words:\n",
    "            for i in range(len(word)):\n",
    "                prefix = word[:i]\n",
    "                if prefix not in prefix_to_words:\n",
    "                    prefix_to_words[prefix] = set([word])\n",
    "                else:\n",
    "                    prefix_to_words[prefix].add(word)\n",
    "        return prefix_to_words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Trie(object):\n",
    "    def __init__(self): \n",
    "        self.words = []\n",
    "        self.childs = collections.defaultdict(Trie)\n",
    "\n",
    "class Solution:\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "\n",
    "        root = Trie()\n",
    "        for word in words:\n",
    "            cur = root\n",
    "            for ch in word:                \n",
    "                cur = cur.childs[ch]\n",
    "                cur.words.append(word)\n",
    "        \n",
    "        def findWordsWithPrefix(prefix: str) -> List[str]:\n",
    "            cur, found = root, True\n",
    "            for ch in prefix:\n",
    "                if ch not in cur.childs:\n",
    "                    found = False\n",
    "                    break\n",
    "                cur = cur.childs[ch]\n",
    "            return cur.words if found else []\n",
    "\n",
    "        res = []\n",
    "        def backtrack(square):\n",
    "            if len(square) == len(square[0]): \n",
    "                res.append(square[:])\n",
    "                return\n",
    "\n",
    "            columnInd = len(square)\n",
    "            prefix = [row[columnInd] for row in square]\n",
    "\n",
    "            for nxt in findWordsWithPrefix(prefix):\n",
    "                square.append(nxt)\n",
    "                backtrack(square)\n",
    "                square.pop()\n",
    "\n",
    "        for word in words:\n",
    "            backtrack([word])\n",
    "        return res \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "        trie = Trie()\n",
    "        for word in words:\n",
    "            trie.insert(word)\n",
    "\n",
    "        result = []\n",
    "        path = []\n",
    "        visited = set()\n",
    "        for word in words:\n",
    "            self.dfs(word, result, path, visited, trie)\n",
    "        return result\n",
    "\n",
    "    def dfs(self, word, result, path, visited, trie):\n",
    "        path.append(word)\n",
    "        if len(path) == len(path[0]):\n",
    "            result.append(path.copy())\n",
    "            path.pop()\n",
    "            return\n",
    "\n",
    "        index = len(path)\n",
    "        for i in range(index, len(word)):\n",
    "            prefix = \"\".join([path[j][i] for j in range(index)])\n",
    "            if not trie.return_words(prefix):\n",
    "                path.pop()        \n",
    "                return\n",
    "\n",
    "        prefix = \"\".join([path[i][index] for i in range(index)])\n",
    "        for next_word in trie.return_words(prefix):\n",
    "            self.dfs(next_word, result, path, visited, trie)\n",
    "        path.pop()\n",
    "\n",
    "class Trie:\n",
    "    def __init__(self):\n",
    "        self.children = {}\n",
    "        self.words = []\n",
    "\n",
    "    def insert(self, word):\n",
    "        node = self\n",
    "        for c in word:\n",
    "            if c not in node.children:\n",
    "                node.children[c] = Trie()\n",
    "            node.words.append(word)\n",
    "            node = node.children[c]\n",
    "\n",
    "    def search(self, word):\n",
    "        node = self\n",
    "        for c in word:\n",
    "            if c not in node.children:\n",
    "                return \n",
    "            node = node.children[c]\n",
    "        return node\n",
    "\n",
    "    def return_words(self, word):\n",
    "        node = self.search(word)\n",
    "        if not node:\n",
    "            return []\n",
    "        return node.words \n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "\n",
    "        n = len(words[0])\n",
    "        ans = []\n",
    "        path = []\n",
    "        d = defaultdict(list)\n",
    "        for x in words:\n",
    "            for i in range(len(x)):\n",
    "                d[x[:i + 1]].append(x)\n",
    "\n",
    "        def dfs(i):\n",
    "            if i == n:\n",
    "                ans.append(path[:])\n",
    "                return\n",
    "            next_prefix = ''.join([x[i] for x in path])\n",
    "\n",
    "            for x in d[next_prefix]:\n",
    "                path.append(x)\n",
    "                dfs(i + 1)\n",
    "                path.pop()\n",
    "        \n",
    "\n",
    "        for x in words:\n",
    "            path.append(x)\n",
    "            dfs(1)\n",
    "            path.pop()\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "from collections import defaultdict, deque\n",
    "\n",
    "class Solution:\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "        \"\"\"\n",
    "\n",
    "        :param words:\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        ret_all = []\n",
    "        len1 = []\n",
    "        len2 = []\n",
    "        len3 = []\n",
    "        len4 = []\n",
    "        for i in words:\n",
    "            if len(i) == 1:\n",
    "                len1.append(i)\n",
    "            elif len(i) == 2:\n",
    "                len2.append(i)\n",
    "            elif len(i) == 3:\n",
    "                len3.append(i)\n",
    "            elif len(i) == 4:\n",
    "                len4.append(i)\n",
    "        for i in len1:\n",
    "            ret_all.append([i])\n",
    "\n",
    "\n",
    "        def process_len(ll, idx_len):\n",
    "            hashM = defaultdict(list)\n",
    "            for i in ll:\n",
    "                for j in range(1, idx_len):\n",
    "                    hashM[i[:j]].append(i)\n",
    "            ret = []\n",
    "\n",
    "            def get_matrix(prev):\n",
    "                if len(prev) == idx_len:\n",
    "                    ret.append(prev)\n",
    "                    return\n",
    "                if len(prev) == 0:\n",
    "                    for i in ll:\n",
    "                        get_matrix([i])\n",
    "                if len(prev) == 1:\n",
    "                    for i in hashM[prev[0][1]]:\n",
    "                        get_matrix(prev + [i])\n",
    "                if len(prev) == 2:\n",
    "                    for i in hashM[prev[0][2] + prev[1][2]]:\n",
    "                        get_matrix(prev + [i])\n",
    "                if len(prev) == 3:\n",
    "                    for i in hashM[prev[0][3] + prev[1][3] + prev[2][3]]:\n",
    "                        get_matrix(prev + [i])\n",
    "\n",
    "            get_matrix([])\n",
    "            return ret\n",
    "\n",
    "\n",
    "        for i in range(2, 5):\n",
    "            if i == 2:\n",
    "                ret_all.extend(process_len(len2, i))\n",
    "            elif i == 3:\n",
    "                ret_all.extend(process_len(len3, i))\n",
    "            elif i == 4:\n",
    "                ret_all.extend(process_len(len4, i))\n",
    "\n",
    "        return [i for i in ret_all if i != []]\n",
    "\n",
    "\n",
    "a = Solution()\n",
    "print(a.wordSquares([\"a\"]))\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "\n",
    "        n = len(words[0])\n",
    "        ans = []\n",
    "        path = []\n",
    "        d = defaultdict(list)\n",
    "        for x in words:\n",
    "            for i in range(len(x)):\n",
    "                d[x[:i]].append(x)\n",
    "\n",
    "        def dfs(i):\n",
    "            if i == n:\n",
    "                ans.append(path[:])\n",
    "                return\n",
    "            next_prefix = ''.join([x[i] for x in path])\n",
    "\n",
    "            for x in d[next_prefix]:\n",
    "                path.append(x)\n",
    "                dfs(i + 1)\n",
    "                path.pop()\n",
    "        \n",
    "\n",
    "        dfs(0)\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Trie(object):\n",
    "    def __init__(self): \n",
    "        self.words = []\n",
    "        self.childs = collections.defaultdict(Trie)\n",
    "\n",
    "class Solution:\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\n",
    "\n",
    "        root = Trie()\n",
    "        for word in words:\n",
    "            cur = root\n",
    "            for ch in word:                \n",
    "                cur = cur.childs[ch]\n",
    "                cur.words.append(word)\n",
    "        \n",
    "        def findWordsInPreifx(prefix: str) -> List[str]:\n",
    "            cur = root\n",
    "            for ch in prefix:\n",
    "                cur = cur.childs[ch]\n",
    "            return cur.words\n",
    "\n",
    "        res = []\n",
    "        def backtrack(square):\n",
    "            if len(square) == len(square[0]): \n",
    "                res.append(square[:])\n",
    "                return\n",
    "\n",
    "            columnInd = len(square)\n",
    "            prefix = [row[columnInd] for row in square]\n",
    "\n",
    "            for nxt in findWordsInPreifx(prefix):\n",
    "                square.append(nxt)\n",
    "                backtrack(square)\n",
    "                square.pop()\n",
    "\n",
    "        for word in words:\n",
    "            backtrack([word])\n",
    "        return res \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Trie:\r\n",
    "    def __init__(self):\r\n",
    "        self.child = dict()\r\n",
    "        self.isWord = False\r\n",
    "\r\n",
    "    def insert(self, word: str) -> None:    \r\n",
    "        root = self\r\n",
    "        for c in word:\r\n",
    "            if c not in root.child:\r\n",
    "                root.child[c] = Trie()\r\n",
    "            root = root.child[c]\r\n",
    "        root.isWord = True\r\n",
    "    \r\n",
    "    def starts_with(self, prefix: str) -> bool:\r\n",
    "        root = self\r\n",
    "        for c in prefix:\r\n",
    "            if c not in root.child:\r\n",
    "                return False\r\n",
    "            root = root.child[c]\r\n",
    "        return True\r\n",
    "    \r\n",
    "    def is_word(self, word: str) -> bool:\r\n",
    "        root = self\r\n",
    "        for c in word:\r\n",
    "            if c not in root.child:\r\n",
    "                return False\r\n",
    "            root = root.child[c]\r\n",
    "        return root.isWord\r\n",
    "\r\n",
    "    def query(self, prefix: str) -> List[str]:\r\n",
    "        res = []\r\n",
    "        root = self\r\n",
    "        for c in prefix:\r\n",
    "            if c not in root.child:\r\n",
    "                return []\r\n",
    "            root = root.child[c]\r\n",
    "\r\n",
    "        def dfs(prefix: str,  root) -> None:\r\n",
    "            if root.isWord == True:\r\n",
    "                res.append(prefix)\r\n",
    "            for c in root.child:\r\n",
    "                dfs(prefix + c, root.child[c])\r\n",
    "\r\n",
    "        dfs(prefix, root)\r\n",
    "        return res\r\n",
    "\r\n",
    "class Solution:\r\n",
    "    def wordSquares(self, words: List[str]) -> List[List[str]]:\r\n",
    "        #挑选单词时顺序是随意的\r\n",
    "        #关于主对角线对称 当前已有的行们决定下一行的前缀\r\n",
    "\r\n",
    "        #------ 先建trie树\r\n",
    "        T = Trie()\r\n",
    "        for word in words:\r\n",
    "            T.insert(word)\r\n",
    "        \r\n",
    "        #------------------- 回溯\r\n",
    "        def backtrace(idx: int) -> None:\r\n",
    "            if idx == wordLen:          #剪枝\r\n",
    "                res.append(path[:])\r\n",
    "                return\r\n",
    "            nxt_prefix = \"\"             #下一行的前缀\r\n",
    "            for i in range(idx):\r\n",
    "                nxt_prefix += path[i][idx]\r\n",
    "\r\n",
    "            nxt_words_with_prefix = T.query(nxt_prefix)\r\n",
    "            for word in nxt_words_with_prefix:\r\n",
    "                path.append(word)\r\n",
    "                backtrace(idx + 1)\r\n",
    "                path.pop(-1)\r\n",
    "\r\n",
    "        #---------------- 运行回溯算法\r\n",
    "        wordLen = len(words[0])\r\n",
    "        res = []\r\n",
    "\r\n",
    "        path = []\r\n",
    "        for word in words:\r\n",
    "            path.append(word)\r\n",
    "            backtrace(1)\r\n",
    "            path.pop(-1)\r\n",
    "        return res"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
